import logging
from langchain_core.prompts import HumanMessagePromptTemplate
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.runnables.base import RunnableLambda
from langchain_core.output_parsers import StrOutputParser
from rag.chroma_conn import ChromaDB
from package.SysConfigManager import SysConfigManager
from langchain.retrievers import ContextualCompressionRetriever

# 配置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
SysConfigManager = SysConfigManager()
persist_path =  SysConfigManager.get_config_path("CHROMADB_PERSIST_DIRECTORY")

class RagManager:
    def __init__(self,
                 chroma_server_type="http",
                 host="localhost", port=8000,
                 persist_path=persist_path,
                 llm=None, embed=None, rerank=None):
        self.llm = llm
        self.embed = embed
        self.rerank = rerank

        chrom_db = ChromaDB(chroma_server_type=chroma_server_type,
                            host=host, port=port,
                            persist_path=persist_path,
                            embed=embed)
        self.store = chrom_db.get_store()

    def get_chain(self):
        """获取RAG查询链"""

        user_prompt = HumanMessagePromptTemplate.from_template(template="""
        请根据用户从私有知识库检索出来的上下文来回答用户的问题！
        请注意：
            1，如果用户的问题不在上下文中，请直接回答不知道！
            2，不要做任何解释，直接输出最终的结果即可！
        检索出的上下文为：
        {context}
        用户的问题为：
        {question}
        答案为：
        """)
        # RAG系统经典的 Prompt (A 增强的过程)
        prompt = ChatPromptTemplate.from_messages(messages=[user_prompt])
        # 将 format_docs 方法包装为 Runnable
        # format_docs_runnable = RunnableLambda(self.format_docs)
        # RAG 链
        rag_chain = (prompt
                | self.llm
                | StrOutputParser()
        )

        return rag_chain

    def format_docs(self, docs):
        # 返回检索到的资料文件名称
        logging.info(f"检索到资料文件个数：{len(docs)}")
        retrieved_files = "\n".join([doc.metadata["source"] for doc in docs])
        logging.info(f"资料文件分别是:\n{retrieved_files}")

        retrieved_content = "\n\n".join(doc.page_content for doc in docs)
        logging.info(f"检索到的资料为:\n{retrieved_content}")
        logging.info(f"==="*20)

        return retrieved_content

    def get_retriever(self, k=4, mutuality=0.3):
        retriever = self.store.as_retriever(search_type="similarity_score_threshold",
                                            search_kwargs={"k": k, "score_threshold": mutuality})
        return retriever
    
    # def get_retrieve_result(self, question, k=8, mutuality=0.2, top_n=4):
    #     """获取chroma数据库的查询结果"""
    #     # 粗排
    #     docs = self.store.similarity_search_with_relevance_scores(query=question, 
    #                                                                 k = k, 
    #                                                                 score_threshold=mutuality)
    #     docs = [doc[0].page_content for doc in docs]
    #     # 精排
    #     docs = self.rerank.rerank(documents=docs, query=question, top_n=top_n)
    #     # 拼接起来
    #     context = "\n\n".join(docs)

    #     return context

    def get_chain_result(self, question, k=10, mutuality=0.3):
        """获取RAG查询结果"""
        retriever = self.get_retriever(k, mutuality)
        compression_retriever = ContextualCompressionRetriever(
            base_compressor=self.rerank, base_retriever=retriever
        )
        compressed_docs = compression_retriever.invoke(question)
        self.format_docs(compressed_docs)
        context = "\n\n".join([doc.page_content for doc in compressed_docs])
        rag_chain = self.get_chain()

        return rag_chain.invoke(input={"context": context, "question": question})